import json
import struct
import numpy as np
import serial
from flask import Flask, render_template, jsonify
import matplotlib.pyplot as plt
from scipy.optimize import least_squares
import threading
import queue
import os

app = Flask(__name__)

# 确保static文件夹存在
if not os.path.exists('static'):
    os.makedirs('static')

# 基站位置
base_location = [(0, 0), (1, 0), (1, 1)]

# 创建全局队列和锁
data_queue = queue.Queue()
serial_lock = threading.Lock()
ser = None

# 初始化串口
def init_serial():
    global ser
    try:
        ser = serial.Serial('COM5', 115200, timeout=1)
        print("串口已成功打开")
        return True
    except serial.SerialException as e:
        print(f"无法打开串口: {e}")
        return False

# 定义误差计算函数
def residuals(p, base_stations, distances):
    x, y = p
    return [np.sqrt((x - bx) ** 2 + (y - by) ** 2) - d for (bx, by), d in zip(base_stations, distances)]

# 定位计算函数
def locate_tag_numerical(base_stations, distances):
    initial_guess = [0, 0]
    result = least_squares(residuals, initial_guess, args=(base_stations, distances))
    return (result.x[0], result.x[1]) if result.success else None

# 数据包解析函数
def parse_data_packet(packet_data):
    packet_data = packet_data[:-2]  # 去除末尾两个字节
    try:
        unpacked_data = struct.unpack('<4B 10H 2B', packet_data)
        distances = [unpacked_data[5] / 100, unpacked_data[6] / 100, unpacked_data[7] / 100]
        loc = locate_tag_numerical(base_location, distances)

        result = {
            'tag_address': unpacked_data[3],
            'frame_num': unpacked_data[4],
            'distance1': unpacked_data[5],
            'distance2': unpacked_data[6],
            'distance3': unpacked_data[7],
            'accel_x': unpacked_data[8] / 100.0,
            'accel_y': unpacked_data[9] / 100.0,
            'accel_z': unpacked_data[10] / 100.0,
            'roll': unpacked_data[11] / 100.0,
            'pitch': unpacked_data[12] / 100.0,
            'yaw': unpacked_data[13] / 100.0,
            'position': loc
        }
        return result
    except struct.error as e:
        print(f"数据包解析错误: {e}")
        return None

# 保存解析结果到 data.json
def save_to_json(data):
    try:
        with open('data.json', 'a') as f:
            json.dump(data, f)
            f.write('\n')
    except IOError as e:
        print(f"保存数据错误: {e}")

# 绘图函数# 绘图函数
def plot_location(data_list):
    plt.ioff()  # 关闭交互模式
    fig, ax = plt.subplots(figsize=(8, 8))  # 设置图表大小
    
    # 设置图表基本属性
    ax.set_xlim(-1.5, 1.5)
    ax.set_ylim(-1.5, 1.5)
    ax.set_xlabel('X position', fontsize=12)
    ax.set_ylabel('Y position', fontsize=12)
    ax.set_title('UWB Tag Position', fontsize=16, fontweight='bold')
    
    # 绘制基站位置
    colors = ['blue', 'green', 'red']
    for i, (bx, by) in enumerate(base_location):
        ax.scatter(bx, by, color=colors[i], s=200, edgecolor='black', linewidth=1.5, label=f'Base {i + 1}')
        ax.text(bx, by + 0.05, f'Base {i + 1}', ha='center', va='bottom', fontsize=12, color='black')
    
    # 绘制最新位置
    if data_list:
        latest_data = data_list[-1]
        if latest_data['position'] is not None:
            x, y = latest_data['position']
            ax.scatter(x, y, color='orange', s=200, edgecolor='black', linewidth=1.5, label='Current Tag Position')
            ax.text(x, y + 0.05, 'Tag', ha='center', va='bottom', fontsize=12, color='black')
    
    # 添加最新的加速度和位置信息
    if data_list:
        latest_data = data_list[-1]
        if latest_data['position'] is not None:
            x, y = latest_data['position']
            accel_text = f'Accel: X={latest_data["accel_x"]:.2f}, Y={latest_data["accel_y"]:.2f}, Z={latest_data["accel_z"]:.2f}'
            pos_text = f'Position: X={x:.2f}, Y={y:.2f}'
            ax.text(0.05, 0.95, accel_text, transform=ax.transAxes, fontsize=12, verticalalignment='top', bbox=dict(facecolor='white', alpha=0.5))
            ax.text(0.05, 0.90, pos_text, transform=ax.transAxes, fontsize=12, verticalalignment='top', bbox=dict(facecolor='white', alpha=0.5))
    
    ax.legend(loc='upper right', fontsize=12)
    ax.grid(True, linestyle='--', alpha=0.7)
    
    # 保存图片
    try:
        plt.savefig('static/location_plot.png', bbox_inches='tight', dpi=300)
    except Exception as e:
        print(f"保存图片错误: {e}")
    finally:
        plt.close()

# 串口读取线程函数
def serial_reader():
    global ser
    if not init_serial():
        return

    try:
        while True:
            with serial_lock:
                if ser and ser.in_waiting >= 28:
                    packet_data = ser.read(28)
                    parsed_data = parse_data_packet(packet_data)
                    if parsed_data and parsed_data['position'] is not None:
                        data_queue.put(parsed_data)
                        save_to_json(parsed_data)
    except Exception as e:
        print(f"串口读取错误: {e}")
    finally:
        if ser:
            with serial_lock:
                ser.close()
                ser = None

# 数据读取函数
def read_serial_data():
    data_list = []
    try:
        while len(data_list) < 10:
            try:
                data = data_queue.get_nowait()
                data_list.append(data)
            except queue.Empty:
                break
    except Exception as e:
        print(f"数据读取错误: {e}")
    return data_list

# Flask路由
@app.route('/')
def index():
    data_list = read_serial_data()
    if data_list:
        plot_location(data_list)
    return render_template('index.html')

@app.route('/latest_data')
def latest_data():
    data_list = read_serial_data()
    if data_list:
        latest_data = data_list[-1]
        return jsonify(latest_data)
    else:
        return jsonify({"message": "No data available"}), 404

if __name__ == '__main__':
    # 启动串口读取线程
    serial_thread = threading.Thread(target=serial_reader, daemon=True)
    serial_thread.start()
    
    # 启动Flask应用
    app.run(debug=True, use_reloader=False)